# coding = utf-8

'''
remove small tumor, size < 10
'''

import os
import numpy as np
import cv2
from scipy import ndimage
from skimage import measure
import prettytable as pt
from pathlib2 import Path

def remove_small_tumor():
    data_path = "/datasets/3DIRCADB/chengkun"
    save_path = "/datasets/3DIRCADB/chengkun_remove"

    for i in range(20):
        total_tumor = 1
        case_id = "case_{}".format(str(i).zfill(5))
        case_path = os.path.join(data_path, case_id)
        segmentation_path = os.path.join(case_path, "segmentation")
        image_path = os.path.join(case_path, "imaging")

        segmentation = []
        image = []
        for item in sorted(os.listdir(segmentation_path)):
            item_file = os.path.join(segmentation_path, item)
            data = np.load(item_file)
            segmentation.append(data)

        for item in sorted(os.listdir(image_path)):
            item_file = os.path.join(image_path, item)
            data = np.load(item_file)
            image.append(data)

        segmentation = np.array(segmentation)
        image = np.array(image)

        tumor = np.zeros(segmentation.shape)
        tumor[segmentation == 2] = 1

        liver_before = (segmentation == 1).sum()
        tumor_before = (segmentation == 2).sum()

        print("*" * 35, case_id, "*" * 35)
        [tumor_labels, num] = measure.label(tumor, return_num=True)
        region = measure.regionprops(tumor_labels)

        tb = pt.PrettyTable()
        tb.field_names = ["tumor NO.", "tumor size", "slice begin", "slice end", "total slice", "avg per slice"]
        remove_list = []
        for j in range(num):
            avg_slice = int(region[j].area / (region[j].bbox[3] - region[j].bbox[0]))
            tb.add_row([
                j, region[j].area, region[j].bbox[0], region[j].bbox[3], region[j].bbox[3] - region[j].bbox[0],
                int(region[j].area / (region[j].bbox[3] - region[j].bbox[0]))]
            )


            if avg_slice <= 10:
                for [x,y,z] in region[j].coords:
                    segmentation[x,y,z] = 1
                remove_list.append(j)
            else:
                total_tumor += 1
                for [x,y,z] in region[j].coords:
                    segmentation[x,y,z] = total_tumor

        liver_after = (segmentation== 1).sum()
        tumor_after = (segmentation >= 2).sum()

        print(tb)
        print("remove tumor :", remove_list)
        print("liver before : ", liver_before, ", liver after : ", liver_after)
        print("tumor before : ", tumor_before, ", tumor after : ", tumor_after)
        print("label : ", np.unique(segmentation))

        destination_patient_path = os.path.join(save_path, "case_{}".format(str(i).zfill(5)))
        destination_patient_path = Path(destination_patient_path)
        destination_image_path = destination_patient_path / "imaging"
        destination_mask_path = destination_patient_path / "segmentation"
        print(destination_image_path, destination_mask_path)
        if not destination_image_path.exists():
            destination_image_path.mkdir(parents=True)
        if not destination_mask_path.exists():
            destination_mask_path.mkdir(parents=True)

        for j in range(segmentation.shape[0]):
            np.save(os.path.join(str(destination_image_path), "{}.npy".format(str(j).zfill(3))), image[j])
            np.save(os.path.join(str(destination_mask_path), "{}.npy".format(str(j).zfill(3))), segmentation[j])

        print()

if __name__ == '__main__':
    remove_small_tumor()

